import torch
import torch.nn as nn
import torch.optim as optim
from data import BrainMRIDataset

# 假设 UNet 和 BrainMRIDataset 已经定义好
from models import UNet
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import DataLoader, random_split
from torchvision import transforms


def train(model, device, train_loader, optimizer, criterion, epoch):
    model.train()
    running_loss = 0.0
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        if batch_idx % 10 == 0:
            print(
                f"Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}"
            )
    return running_loss / len(train_loader)


def validate(model, device, val_loader, criterion):
    model.eval()
    val_loss = 0
    with torch.no_grad():
        for data, target in val_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            val_loss += criterion(output, target).item()
    val_loss /= len(val_loader)
    print(f"\nValidation set: Average loss: {val_loss:.4f}\n")
    return val_loss


# 基本参数
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
batch_size = 4
epochs = 10
lr = 1e-3
gamma = 0.7

# 数据加载和预处理
transform = transforms.Compose(
    [transforms.ToTensor(), transforms.Normalize(mean=[0.5], std=[0.5])]
)

# 加载自定义数据集
data_dir = "lgg-mri-segmentation/kaggle_3m"  # 确保此路径正确
full_dataset = BrainMRIDataset(base_dir=data_dir, transform=transform)
train_size = int(0.8 * len(full_dataset))
val_size = len(full_dataset) - train_size
train_dataset, val_dataset = random_split(full_dataset, [train_size, val_size])

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

# 初始化模型
model = UNet().to(device)

# 损失函数和优化器
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
scheduler = StepLR(optimizer, step_size=1, gamma=gamma)

# 训练和验证循环
for epoch in range(1, epochs + 1):
    train_loss = train(model, device, train_loader, optimizer, criterion, epoch)
    val_loss = validate(model, device, val_loader, criterion)
    scheduler.step()
    print(f"Epoch {epoch}: Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")

# 保存模型
torch.save(model.state_dict(), "unet_model_brain_mri.pth")
